package com.lry.util.mysql;

import com.alibaba.excel.context.AnalysisContext;
import com.alibaba.excel.event.AnalysisEventListener;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.function.Consumer;
import java.util.stream.Collectors;
import lombok.EqualsAndHashCode;

/**
 * Excel读取监听器
 * 默认head为第0行
 *
 * @author ALI
 * @since 2021/12/13
 */
@EqualsAndHashCode(callSuper = true)
public class ReadExcelListener extends AnalysisEventListener<Map<Integer, String>> {

    public static final int MAX_LINE = 10000;
    public static final int LATCH_COUNT = 500;
    private static final String N = "\r\n";
    private final Map<Integer, String> headerMap = new HashMap<>();
    private final Map<Integer, Integer> headerLength = new HashMap<>();
    private final CountDownLatch latch = new CountDownLatch(5000);
    private final Consumer<Map<Integer, String>> consumer;
    private final ExecutorService pool;
    private final Consumer<ByteArrayOutputStream> bodyConsumer;
    private ByteArrayOutputStream os = new ByteArrayOutputStream();
    private boolean header;
    private int lines = 0;
    private long totalLines = 0;
    private int threadCount = 0;

    public ReadExcelListener(Consumer<Map<Integer, String>> headerConsumer,
                             Consumer<ByteArrayOutputStream> bodyConsumer) {
        this.consumer = headerConsumer;
        this.bodyConsumer = bodyConsumer;
        pool = new ThreadPoolExecutor(4, 6,
                                      0L, TimeUnit.MILLISECONDS,
                                      new LinkedBlockingQueue<Runnable>(200));
    }

    private static String toStr(String d) {
        if (d == null) {
            return "";
        } else {
            if (d.contains(",")) {
                return "\"" + d + "\"";
            }
            return d;
        }
    }

    @Override
    public void invoke(Map<Integer, String> integerStringMap, AnalysisContext analysisContext) {
        try {
            ++totalLines;
            if (!header) {
                header = true;
                headerMap.putAll(integerStringMap);
                Optional.ofNullable(consumer).ifPresent(d -> d.accept(integerStringMap));
            } else {
                if (++lines > MAX_LINE) {
                    os.flush();
                    ByteArrayOutputStream tempOs = os;
                    syncExecute(tempOs);
                    os = new ByteArrayOutputStream();
                    lines = 0;
                }
                List<String> collect = integerStringMap.values().stream().map(ReadExcelListener::toStr)
                                                       .collect(Collectors.toList());
                String t = String.join(",", collect) + N;
                os.write(t.getBytes(StandardCharsets.UTF_8));
            }
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void doAfterAllAnalysed(AnalysisContext analysisContext) {
        try {
            os.flush();
            syncExecute(os);
            os.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取总行数
     */
    public long getTotalLines() {
        return totalLines;
    }

    /**
     * 等待所有任务执行完成
     */
    public void waitThread() {
        while (true) {
            try {
                TimeUnit.MILLISECONDS.sleep(20);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            if ((LATCH_COUNT - latch.getCount()) == threadCount) {
                return;
            }
        }
    }

    /**
     * 表头
     */
    public Map<String, Integer> getHeader() {
        Map<String, Integer> result = new HashMap<>();
        for (Map.Entry<Integer, String> entry : headerMap.entrySet()) {
            result.put(entry.getValue(), headerLength.get(entry.getKey()));
        }
        return result;
    }

    private void syncExecute(ByteArrayOutputStream tempOs) {
        ++threadCount;
        pool.execute(new Thread(() -> {
            bodyConsumer.accept(tempOs);
        }));
    }

}
